from contextlib import contextmanager

from flask_sqlalchemy import SQLAlchemy as _SQLAlchemy


class SQLAlchemy(_SQLAlchemy):

    @contextmanager
    def auto_commit(self):
        try:
            yield
            self.session.commit()
            print("Transaction committed")
        except Exception as e:
            self.session.rollback()
            print(f"Transaction rolled back due to: {e}")
            raise e
